import numpy as np
import pyopencl as cl
from openclinfo import *

class Verlet:
    cl_files = {'verlet':'verlet.cl'}

    def __init__(self,name=None,context=None,flags=''):
        self.name = name #'pair_neighbor_list'
        cl_file = Verlet.cl_files[name] #'neighbor_list.cl'
        cl_code = file(cl_file,"r").read()
        self.prg = cl.Program(context,cl_code).build(options=flags)
        #

    def nve(self,que,cfg,neigh,pot,nstep,dt):
        sglb = (GLOBAL_SIZE,)
        sloc = (LOCAL_SIZE,)
        for i in range(nstep):
            #first half step of v:
            self.prg.step_v(que,sglb,sloc,
                    np.int32(cfg.n_atoms),
                    np.float32(dt*0.5),
                    cfg._f_e,cfg._v_w,cfg._v_w)
            #full step of r:
            self.prg.step_r(que,sglb,sloc,
                    np.int32(cfg.n_atoms),
                    np.float32(dt),
                    cfg._v_w,cfg._r_t,cfg._r_t)
            #cal new f:
            pot.cal_f_e_vir_all(que,cfg,neigh)
            #second half step of v:
            self.prg.step_v(que,sglb,sloc,
                    np.int32(cfg.n_atoms),
                    np.float32(dt*0.5),
                    cfg._f_e,cfg._v_w,cfg._v_w)
